# -*- coding: utf-8 -*-  
'''
LDA模型相关零件

Created on 2021年8月20日
@author: luoyi
'''
import tensorflow as tf

import utils.conf as conf


#    Gibbs层
class GibbsLayer(tf.keras.layers.Layer):
    '''gibbs采样层
        初始化：
            - 主题数：K
            - 总词数：W
            - 两个Dirichlet分布的超参：α，β
            - 每个主题下、每个词出现的计数：wn_topic_word [K, W]
            - 每个主题下的词计数：wn_topic [K]
        前向过程：
            - 初始化batch篇文档，每篇文档下、每个主题的词计数：wn_doc_topic [batch_size, K]
            - 初始化batch篇文档，每篇文档下的词计数：wn_doc [batch_size]
            - 初始化batch篇文档，每篇文档下、每个词的主题：z [batch_size, None(文档实际词数)]
            - 遍历每篇文档doc，每个词word：
                - 计算该词主题分布（从属于每个主题的概率）
                    part1 = (wn_doc_topic[k] + α[k]) / ∑(K) (wn_doc_topic[k] + α[k])
                    part2 = (wn_topic_word[k,w] + β[w]) / ∑(W)(wn_topic_word[k,w] + β[w])
                    p[k] = part1 * part2
                - 从p分布中采样词的主题new_topic
                - 更新相关计数：
                    wn_topic_word[new_topic, word] += 1
                    wn_doc_topic[doc, new_topic] += 1
            - 统计频率，用频率≈概率
                φ[k] = {wn_topic_word[k,w] / ∑(W)wn_topic_word[k,w]}
    '''
    def __init__(self,
                 name='gibbs_layer',
                 batch_size=conf.DATASET_SOHU_NEWS.get_batch_size(),
                 K=conf.LDA.get_k(),                                #    总主题数
                 W=None,                                            #    语料库总词数
                 max_doc_words=conf.LDA.get_max_doc_words(),        #    每篇文档最大词数
                 batch_words=conf.LDA.get_batch_words(),            #    批量计算的词数量（加快计算速度，但会降低精度）
                 **kwargs):
        self._batch_size = batch_size
        
        self._K = K
        self._W = W
        self._max_doc_words = max_doc_words
        self._batch_words = batch_words
        
        super(GibbsLayer, self).__init__(name=name, trainable=False, **kwargs)
        pass
    
    def build(self, input_shape):
        #    初始化全局参数
        
        #    初始化Dirichlet(α)分布的超参，Tensor(K, )
        alpha = tf.convert_to_tensor([50. / self._K], dtype=tf.float32)
        alpha = tf.tile(alpha, multiples=[self._K])
        self._alpha = self.add_weight(name='gibbs/dirichlet/alpha', 
                                      shape=(self._K,), 
                                      dtype=tf.float32, 
                                      initializer='zero', trainable=False)
        self._alpha.assign(alpha)
        
        #    初始化Dirichlet(β)分布的超参，Tensor(W, )
        beta = tf.convert_to_tensor([0.02], dtype=tf.float32)
        beta = tf.tile(beta, multiples=[self._W])
        self._beta = self.add_weight(name='gibbs/dirichlet/beta', 
                                     shape=(self._W,), 
                                     dtype=tf.float32, 
                                     initializer='zero', trainable=False)
        self._beta.assign(beta)
        
        #    初始化每个主题每个词的 词计数    Tensor(K, W)
        self._wn_topic_word = self.add_weight(name='gibbs/wn_topic_word', 
                                              shape=(self._K, self._W), 
                                              dtype=tf.float32, 
                                              initializer='zero', 
                                              trainable=False)
        self._wn_topic_word.assign(tf.ones(shape=(self._K, self._W)) * 0.00001)
        
        #    初始化当前主题
        self._wn_docs_topic = self.add_weight(name='gibbs/wn_docs_topic', 
                                              shape=(self._batch_size, self._K), 
                                              dtype=tf.float32, 
                                              initializer='zero', 
                                              trainable=False)
        self._wn_docs_topic.assign(tf.ones(shape=(self._batch_size, self._K)) * 0.0001)
        
        #    主题-词分布
        self._topic_word_prob = self.add_weight(name='gibbs/topic_word_prob', 
                                                shape=(self._K, self._W), 
                                                dtype=tf.float32, 
                                                initializer='zero', 
                                                trainable=False)
        #    文档-主题分布
        self._docs_topic_prob = self.add_weight(name='gibbs/docs_topic_prob', 
                                                shape=(self._batch_size, self._K), 
                                                dtype=tf.float32, 
                                                initializer='zero', 
                                                trainable=False)
        
        super(GibbsLayer, self).build(input_shape)
        pass
    
    #    取当前主题-词 词计数
    def wn_topic_word(self):
        return self._wn_topic_word
    #    取当前文档-主题 词计数
    def wn_docs_topic(self):
        return self._wn_docs_topic 
    
    #    取当前主题-词分布
    def topic_word_prob(self):
        return self._topic_word_prob
    #    取当前文档-主题分布
    def docs_topic_prob(self):
        return self._docs_topic_prob
    
    
    def call(self, inputs, training=None):
        '''
            @param inputs: Tensor(batch_size, max_doc_words)    -1表示掩码
        '''
        #    按batch_words切分数据块
        X = tf.reshape(inputs, shape=(-1, self._max_doc_words // self._batch_words, self._batch_words))
        
        #    初始化文档-主题 词计数    Tensor(batch_size, K)
        wn_docs_topic = tf.zeros(shape=(self._batch_size, self._K), dtype=tf.float32)
        
        #    遍历每篇文档
        m = 0
        for doc in X:
            #    遍历每个词块。每个词块 Tensor(batch_words, )
            for block in doc:
                block = tf.cast(block, dtype=tf.int32)
                
                #    如果block中全是-1，则忽略次block
                if (tf.math.reduce_mean(block) == -1): continue
                #    过滤掉-1
                block = tf.gather_nd(block, indices=tf.where(block >= 0))

                #    计算词的主题分布
                p = self.topic_distribution(m, block, self._wn_docs_topic, self._wn_topic_word, self._alpha, self._beta)
                
                #    从此的主题分布中采样新的主题
                new_topic = self.sample_topic_from_p(p)
                
                #    更新计数
                wn_docs_topic, wn_topic_word = self.update_count(m, new_topic, block, self._wn_docs_topic, self._wn_topic_word)
                
                #    训练阶段更新相关计数
                if (training):
                    self._wn_topic_word.assign(wn_topic_word)
                    self._wn_docs_topic.assign(wn_docs_topic)
                    pass
                
                pass
            
            m += 1
            pass
        
        #    计算文档-主题 分布
        self.update_doc_topic_prob(wn_docs_topic)
        #    如果是训练阶段，计算主题-词分布
        if (training): self.update_topic_word_prob(self._wn_topic_word)
        return inputs

    #    计算词的主题分布
    def topic_distribution(self, m, block, wn_docs_topic, wn_topic_word, alpha, beta):
        '''
            part1 = (wn_doc_topic[k] + α[k]) / ∑(K) (wn_doc_topic[k] + α[k])
            part2 = (wn_topic_word[k,w] + β[w]) / ∑(W)(wn_topic_word[k,w] + β[w])
            p[k] = part1 * part2
            
            @param m: 第m篇文档
            @param n: 第n个词
            @param block: 词块，每个值为wid Tensor(batch_words, )
            @param wn_docs_topic: 文档的主题，词计数 Tensor(batch_size, K)
            @param wn_topic_word: 主题的词 词计数 Tensor(K, W)
            @param alpha: 文档-主题的Dirichlet分布参数 Tensor(K, )
            @param beta: 主题-词的Dirichlet分布参数 Tensor(W, )
            @param temp_wn_topic_word: 用于推理的主题的词 词计数 Tensor(K, W)
        '''
        part1 = (wn_docs_topic[m] + alpha) / tf.math.reduce_sum(wn_docs_topic[m] + alpha)    #    Tensor(K, )
        #    计算(wn_topic_word[k,w] + β[w]) / ∑(W)(wn_topic_word[k,w] + β[w])
        #    分子部分
        p2_wn = tf.gather(wn_topic_word, indices=block, axis=-1)                             #    Tensor(K, batch_words)
        p2_beta = tf.gather(beta, indices=block)                                             #    Tensor(batch_words, )
        p2_molecule = p2_wn + p2_beta[tf.newaxis, :]                                         #    Tensor(K, batch_words)
        #    分母部分
        p2_denominator = tf.math.reduce_sum(wn_topic_word + beta[tf.newaxis, :], axis=-1)    #    Tensor(K, )
        part2 = p2_molecule / p2_denominator[:, tf.newaxis]                                  #    Tensor(K, batch_words)
        
        #    计算(wn_doc_topic[k] + α[k]) / ∑(K) (wn_doc_topic[k] + α[k])
        p = part1[:, tf.newaxis] * part2                                                     #    Tensor(K, batch_words)
        p = tf.transpose(p, perm=[1, 0])                                                     #    Tensor(batch_words, K)每行表示词的概率分布
        return p
    #    从主题分布中采样新的主题
    def sample_topic_from_p(self, p):
        '''
            @param p: 每个主题的概率 Tensor(batch_words, K)
            @return: Tensor(batch_words, )
        '''
        return tf.random.categorical(logits=tf.math.log(p), num_samples=1, dtype=tf.int32)[:, 0]
    #    更新相关计数
    def update_count(self, m, new_topic, block, wn_docs_topic, wn_topic_word):
        '''
            @param m: 第m篇文档
            @param n: 第n个词
            @param new_topic: 新的主题索引（相对于0 ~ K的索引） Tensor(batch_words, )
            @param block: 词块，每个值为wid Tensor(batch_words, )
            @param wn_docs_topic: 文档-主题 词计数 Tensor(batch_size, K)
            @param wn_topic_word: 主题的词 词计数 Tensor(K, W)
        '''
        #    更新文档-主题 的 词计数
        res = self.uniqon_count(new_topic)
        st_idx = tf.concat([(tf.ones_like(res[:, 0]) * tf.cast(m, dtype=tf.int64))[:, tf.newaxis], res[:, 0][:, tf.newaxis]], axis=-1)
        st_val = tf.cast(res[:, 1], dtype=tf.float32)
        add_val = tf.SparseTensor(indices=st_idx, values=st_val, dense_shape=wn_docs_topic.shape)
        add_val = tf.sparse.to_dense(add_val)
        wn_docs_topic = wn_docs_topic + add_val
        
        #    更新主题-词 的 词计数
        nt_w = new_topic * self._W + block                                              #    Tensor(K, batch_words)
        res = self.uniqon_count(nt_w)
        nt = res[:, 0] // self._W
        w = res[:, 0] % self._W
        st_idx = tf.concat([nt[:, tf.newaxis], w[:, tf.newaxis]], axis=-1)
        st_val = tf.cast(res[:, 1], dtype=tf.float32)
        add_val = tf.SparseTensor(indices=st_idx, values=st_val, dense_shape=wn_topic_word.shape)
        add_val = tf.sparse.to_dense(add_val)
        wn_topic_word = wn_topic_word + add_val
        return wn_docs_topic, wn_topic_word
    #    计算向量中的元素个数，并返回去重后的向量，和他对应的个数。并按向量值升序
    def uniqon_count(self, vec):
        '''
            @param vec: 待计算的向量
            @param max_val: 向量的最大取值
            @return: Tensor(vec中不重复的元素个数, 2)
                        0: vec中不重复的元素
                        1: 该元素的个数
        '''
        #    通过去重的第二个输出，每个位置表示为相对去重结果的索引。可以用他来统计数量
        vec_uniqon = tf.unique(vec)
        
        #    计算每个元素个数
        vec_idx_oh = tf.one_hot(vec_uniqon[1], depth=tf.math.reduce_max(vec_uniqon[1])+1)
        vec_unqion_count = tf.math.reduce_sum(vec_idx_oh, axis=0)
        #    排序
        vec = tf.concat([tf.cast(vec_uniqon[0], dtype=tf.int64)[:, tf.newaxis],
                        tf.cast(vec_unqion_count, dtype=tf.int64)[:, tf.newaxis]], axis=-1)
        vec = tf.gather(vec, indices=tf.argsort(vec[:, 0]), axis=0)
#         vec = tf.sort(vec, axis=0)
        
        return vec
    
    #    计算当前主题-词分布
    def update_topic_word_prob(self, wn_topic_word):
        '''重新计算当前的主题-词分布
            @param wn_topic_word: 主题-词 的词计数 Tensor(K, W)
        '''
        p = wn_topic_word / tf.math.reduce_sum(wn_topic_word, axis=-1)[:, tf.newaxis]
        self._topic_word_prob.assign(p)
        pass
    #    计算文档-主题分布
    def update_doc_topic_prob(self, wn_docs_topic):
        '''计算当前的文档-主题分布
            @param wn_docs_topic: 文档-主题 词计数 Tensor(batch_size, K)
        '''
        p = wn_docs_topic / tf.math.reduce_sum(wn_docs_topic, axis=-1)[:, tf.newaxis]
        self._docs_topic_prob.assign(p)
        pass
    
    
    pass





